1738E - Balance Addicts - CodeForces Solution

#include <stdio.h>
#include <iostream>
#include <vector>
#include <math.h>
#include <time.h>
#include <algorithm>
#include <string.h>
#include <string>
#include <deque>
#include <ctype.h>
#include <limits.h>
#include <map>
#include <set>
#include <stack>
#include <queue>
#include <functional>
#include <limits.h>
#include <random>
#include <list>
#include <bitset>

using namespace std;
typedef long long ll;

const double PI = acos(-1.0);
const char* YES = "Yes";
const char* NO = "No";

void solve();
mt19937 rng(time(0));

ll qpow(ll x, ll p, ll MOD)
{
  ll res = 1, v = x;
  while (p)
  {
    if (p & 1)
    {
      res *= v;
      res %= MOD;
    }
    v *= v;
    v %= MOD;
    p >>= 1;
  }
  return res;
}
ll inv(ll x, ll MOD)
{
  return qpow(x, MOD - 2, MOD);
}

void get_factors(ll x, vector<int>& v) {
  for(int i = 2; i * i <= x; ++i) {
    if(x % i == 0) {
      v.push_back(i);
      if(i * i != x) {
        v.push_back(x / i);
      }
    }
  }
  v.push_back(x);
}

void get_primefactors(ll x, vector<int>& v) {
  int t = x;
  for(int i = 2; i * i <= x; ++i) {
    if(t % i == 0) {
      v.push_back(i);
      while(t % i == 0) {
        t /= i;
      }
    }
  }
  if(t != 1) {
    v.push_back(t);
  }
}

void init()
{
#ifdef _LOCAL_
  if (!freopen("case.txt", "r", stdin))
  {
    perror("open file failed!:");
    exit(1);
  }
  // if (!freopen("main.ans", "w", stdout))
  // {
  //   perror("open file failed!:");
  //   exit(1);
  // }
#endif
}
#ifdef _LOCAL_
void dbg_out()
{
  cerr << endl;
}
template <typename Head, typename... Tail>
void dbg_out(Head H, Tail... T)
{
  cerr << ' ' << H;
  dbg_out(T...);
}
#define dbg(...) cerr << "(" << #__VA_ARGS__ << "):", dbg_out(__VA_ARGS__)
#else
#define dbg(...)
#endif

ll read()
{
  int ch = getchar();
  if (ch == EOF)
  {
    return INT_MIN;
  }
  while (!isdigit(ch) && ch != '-')
  {
    ch = getchar();
    if (ch == EOF)
    {
      return INT_MIN;
    }
  }
  ll res = 0, sign = 1;
  if (ch == '-')
  {
    ch = getchar();
    sign = -1;
  }
  while (isdigit(ch))
  {
    res = res * 10 + ch - '0';
    ch = getchar();
  }
  return res * sign;
}

int nextchar()
{
  int ch = getchar();
  while (ch == ' ' || ch == '\n' || ch == '\r')
  {
    ch = getchar();
  }
  return ch;
}

void _print(bool x) {x? printf("YES") : printf("NO");}
void _print(double x) {printf("%.12lf", x);}
void _print(char x) { putchar(x);}
void _print(int x) { printf("%d", x);}
void _print(ll x) { 
#ifdef _LOCAL_
  printf("%I64d", x);
#else
  printf("%lld", x);
#endif
}
void _print(unsigned int x) { printf("%u", x);}
void _print(const char* s) { printf("%s", s);}
void _print(string& s) {printf("%s", s.c_str());}
template<typename T> void _print(vector<T>& v) {
  for(int i = 0; i < (int)v.size(); ++i) {
    _print(v[i]);
    if(i + 1 < (int)v.size()) {
      putchar(' ');
    }
  }
}
template<typename T1, typename T2> void _print(pair<T1, T2>& p) {
  _print(p.first);
  putchar(' ');
  _print(p.second);
}


void out() {}
template<typename Head, typename... Args> void out(Head val, Args... rem) {
  _print(val);
  if(sizeof...(rem) > 0) {
    putchar(' ');
  }
  out(rem...);
}

template<typename... Args> void outl(Args... args) {
  out(args...);
  out('\n');
}

template<typename... Args> void outs(Args... args) {
  out(args...);
  out(' ');
}

int main()
{
  init();
  solve();
  return 0;
}
/** 写之前 : 
 *  1. MAXN 是否定义正确？
 *  2. 要把所有数据读完后再计算!
 */
/*********************code start here*********************/
// 常用 2 的倍数
// 1 << 10 = 1024 --- 10^3
// 1 << 17 = 131072 --- 10^5
// 1 << 18 = 262144 --- 2 * 10^5
// 1 << 20 = 1048576 --- 10^6

/** 提交前检查项汇总
 1. 记忆化搜索：不要更改函数 f(x, y, ...) 里的参数 x,y 否则可能会无限递归
    建议采取定义新的变量的方式
 2. 数据范围是否是 10^10? 中间结果是否达到 10^10 - 10^18?
 3. 答案 res 和中间过程 是不是 ll!
 4. 不要用 multiset::count!
 5. 注意，用 map 做计数器时，[] 操作可能导致 size() 增加！
*/

const int MOD = 998244353;
// 注意：组合数必须 mod = 质数 且 n, k < mod 时才可用
class Comb {
    vector<int> Facs, Invs;
    void expand(size_t n) {
        while(Facs.size() < n + 1) Facs.push_back(1ll * Facs.back() * Facs.size() % MOD);
        if(Invs.size() < n + 1) { // 线性求阶乘的逆元
            Invs.resize(n + 1, 0);
            Invs.back() = 1;
            for(int a = Facs[n], p = MOD - 2; p; p >>= 1, a = 1ll * a * a % MOD)
                if(p & 1) Invs.back() = 1ll * Invs.back() * a % MOD; // 快速乘求 n! 的逆元
            for(int j = n-1; !Invs[j]; --j) Invs[j] = 1ll * Invs[j+1] * (j + 1) % MOD;
        }
    }
public:
    Comb() : Facs({1}), Invs({1}) {}
    Comb(int n) : Facs({1}), Invs({1}) { expand(n); }
    int operator() (int n, int k) {
        if(k > n) return 0;
        expand(n);
        return (1ll * Facs[n] * Invs[n-k] % MOD) * Invs[k] % MOD; 
    }
};

Comb comb;

void solve() {
  int t = read();
  while(t--) {
    int n = read();
    int a[n + 1];
    long long s[n + 2];
    for(int i = 1; i <= n; ++i) a[i] = read();
    s[0] = 0;
    for(int i = 1; i <= n; ++i) s[i] = s[i-1] + a[i];
    s[n+1] = s[n];
    int res = 1;
    for(int i = 0, j = n + 1;;) {
      while(j > i && s[n] - s[j-1] < s[i]) --j;
      if(j <= i) {
        break;
      }
      if(s[n] - s[j-1] != s[i]) {
        i++;
        continue;
      }
      if(s[j-1] - s[i] == 0) {
        int d = j-i;
        if(i == 0) d--;
        if(j == n+1) d--;
        res = 1ll * res * qpow(2, d, MOD) % MOD;
        break;
      }
      int p1 = i + 1, p2 = j-1, x = 0, y = 0;
      if(i != 0) x = 1, y = 1;
      while(a[p1] == 0) ++p1, ++x;
      while(a[p2] == 0) --p2, ++y;
      int cur = 0;
      for(int k = 0; k <= min(x, y); ++k) {
        cur += 1ll * comb(x, k) * comb(y, k) % MOD;
        cur %= MOD;
      }
      res = 1ll * res * cur % MOD;
      i = p1;
    }
    outl(res);
  }
}